import numpy as np
import torch

from llmpt.preprocess import Vocab
from llmpt.utils import torch_gather_from_array
from .vocab import VocabForNLP


def make_vocab_transform(vocab_org: Vocab, vocab_dst: Vocab):
    if vocab_org == vocab_dst:
        def transform(indices: int | np.ndarray | torch.Tensor) -> int | np.ndarray | torch.Tensor:
            return indices
    elif isinstance(vocab_org, VocabForNLP):
        if isinstance(vocab_dst, VocabForNLP):
            if vocab_org.base == vocab_dst.base:  # do nothing
                def transform(indices: int | np.ndarray | torch.Tensor) -> int | np.ndarray | torch.Tensor:
                    return indices
            else:
                specials_n = len(VocabForNLP.specials)
                base_transform = make_vocab_transform(vocab_org.base, vocab_dst.base)

                def transform(indices: int | np.ndarray | torch.Tensor) -> int | np.ndarray | torch.Tensor:
                    if isinstance(indices, int):
                        indices = indices if indices < specials_n \
                            else base_transform(indices - specials_n) + specials_n
                    elif isinstance(indices, np.ndarray):
                        indices = np.where(indices < specials_n, indices,
                                           base_transform(indices - specials_n) + specials_n)
                    elif isinstance(indices, torch.Tensor):
                        # noinspection PyTypeChecker
                        indices = torch.where(indices < specials_n, indices,
                                              base_transform(indices - specials_n) + specials_n)
                    else:
                        raise TypeError(f"Unsupported type: {type(indices)}")
                    return indices
        else:
            if vocab_org.base == vocab_dst:  # unwrapping
                def transform(indices: int | np.ndarray | torch.Tensor) -> int | np.ndarray | torch.Tensor:
                    if isinstance(indices, int):
                        indices = indices - len(VocabForNLP.specials)
                        indices = indices if indices >= 0 else -1
                    elif isinstance(indices, np.ndarray):
                        indices = indices - len(VocabForNLP.specials)
                        indices = np.where(indices >= 0, indices, -1)
                    elif isinstance(indices, torch.Tensor):
                        indices = indices - len(VocabForNLP.specials)
                        # noinspection PyTypeChecker
                        indices = torch.where(indices >= 0, indices, -1)
                    else:
                        raise TypeError(f"Unsupported type: {type(indices)}")
                    return indices
            else:
                unwrapping_transform = make_vocab_transform(vocab_org, vocab_org.base)
                base_transform = make_vocab_transform(vocab_org.base, vocab_dst)

                def transform(indices: int | np.ndarray | torch.Tensor) -> int | np.ndarray | torch.Tensor:
                    indices = unwrapping_transform(indices)
                    indices = base_transform(indices)
                    return indices
    else:
        if isinstance(vocab_dst, VocabForNLP):
            if vocab_org == vocab_dst.base:  # wrapping
                def transform(indices: int | np.ndarray | torch.Tensor) -> int | np.ndarray | torch.Tensor:
                    if isinstance(indices, int):
                        indices = (indices + len(VocabForNLP.specials)) if indices >= 0 \
                            else VocabForNLP.UNKNOWN_INDEX
                    elif isinstance(indices, np.ndarray):
                        indices = np.where(indices >= 0, indices + len(VocabForNLP.specials),
                            VocabForNLP.UNKNOWN_INDEX)
                    elif isinstance(indices, torch.Tensor):
                        # noinspection PyTypeChecker
                        indices = torch.where(indices >= 0, indices + len(VocabForNLP.specials),
                            VocabForNLP.UNKNOWN_INDEX)
                    else:
                        raise TypeError(f"Unsupported type: {type(indices)}")
                    return indices
            else:
                base_transform = make_vocab_transform(vocab_org, vocab_dst.base)
                wrapping_transform = make_vocab_transform(vocab_dst.base, vocab_dst)

                def transform(indices: int | np.ndarray | torch.Tensor) -> int | np.ndarray | torch.Tensor:
                    indices = base_transform(indices)
                    indices = wrapping_transform(indices)
                    return indices
        else:
            tokens_org = [vocab_org.get_token(i) for i in range(len(vocab_org))]
            tokens_dst = [vocab_dst.get_token(i) for i in range(len(vocab_dst))]
            if len(tokens_org) <= len(tokens_dst) and tokens_org == tokens_dst[:len(tokens_org)]:  # expanding
                def transform(indices: int | np.ndarray | torch.Tensor) -> int | np.ndarray | torch.Tensor:
                    return indices
            elif len(tokens_org) > len(tokens_dst) and tokens_org[:len(tokens_dst)] == tokens_dst:  # truncating
                dst_size = len(tokens_dst)

                def transform(indices: int | np.ndarray | torch.Tensor) -> int | np.ndarray | torch.Tensor:
                    if isinstance(indices, int):
                        indices = indices if indices < dst_size else -1
                    elif isinstance(indices, np.ndarray):
                        indices = np.where(indices < dst_size, indices, -1)
                    elif isinstance(indices, torch.Tensor):
                        # noinspection PyTypeChecker
                        indices = torch.where(indices < dst_size, indices, -1)
                    else:
                        raise TypeError(f"Unsupported type: {type(indices)}")
                    return indices
            else:  # reindexing
                mapping = []
                for i_org in range(len(vocab_org)):
                    token = vocab_org.get_token(i_org)
                    try:
                        i_dst = vocab_dst.index_token(token)
                    except KeyError:
                        i_dst = -1
                    mapping.append(i_dst)
                mapping = np.asarray(mapping, dtype=np.int64)
                torch_mapping = None

                def transform(indices: int | np.ndarray | torch.Tensor) -> int | np.ndarray | torch.Tensor:
                    if isinstance(indices, int):
                        indices = mapping[indices]
                    elif isinstance(indices, np.ndarray):
                        indices = mapping[indices]
                    elif isinstance(indices, torch.Tensor):
                        nonlocal torch_mapping
                        if torch_mapping is None:
                            torch_mapping = torch.asarray(mapping)
                        device_torch_mapping = torch_mapping.to(indices.device)
                        indices = torch_gather_from_array(device_torch_mapping, indices)
                    else:
                        raise TypeError(f"Unsupported type: {type(indices)}")
                    return indices
    return transform
